using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Text;
using System.Xml;

using Atomic.Libraries;
using Atomic.Libraries.Compression;
using Atomic.Libraries.Mathematics;
using Atomic.Structures;
using Atomic.Thermodynamics.Electrons;

namespace Atomic.Vasp
{
	/// <summary>
	/// VASP output representation. This is the default representation of the newest version of VASP. Other classes are avaiable for
	/// older version and for parsing a reduced set of files.
	/// </summary>
	[Serializable]
	public class VaspResult : BaseVaspResult, IVaspResult
	{
		[NonSerialized] private bool parsed, parsedElectronicStateDensity;
		[NonSerialized] private double energy, magneticMoment, fermiEnergy, electrons, externalPressure;
		[NonSerialized] private Structure initialStructure;
		[NonSerialized] private VaspStructure finalStructure;
		[NonSerialized] private ElectronicStateDensity electronicStateDensity;

		public VaspResult(DirectoryInfo directory)
			: base(directory)
		{
		}

		public static IVaspResult Read()
		{
			// Automatically determines which parser to use.
			throw new NotImplementedException();
		}

		public void Parse()
		{
			if (!parsed)
			{
				parsed = true;

				XmlDocument xml = new XmlDocument();
				xml.LoadXml(Encoding.Default.GetString(Vasprun.GetData()));

				Lattice primitiveLattice;
				int primitiveCells;
				IEnumerable<Atom> atoms;

				ReadOutcar(out primitiveLattice, out primitiveCells, out atoms);
				BuildStructure(xml, primitiveLattice, primitiveCells, atoms);
				ReadPoscar(atoms);
				ReadVasprun(xml);
				ReadOszicar();

				// Optional output files. May be empty.
				//ReadChargeDensity(directory);
				//ReadElectronicStateDensity(directory);
			}
		}

		public void ParseElectronicStateDensity()
		{
			if (!parsedElectronicStateDensity)
			{
				parsedElectronicStateDensity = true;

				// Must do the basic parsing first.
				Parse();

				ReadElectronicStateDensity();
			}
		}

		private void ReadOutcar(out Lattice primitiveLattice, out int primitiveCells, out IEnumerable<Atom> atoms)
		{
			// Atoms as determined by entries in OUTCAR. May be missing in POSCAR/CONTCAR.
			List<string> potcar = new List<string>();
			atoms = new List<Atom>();

			primitiveLattice = null;
			primitiveCells = 0;
			externalPressure = double.NaN;

			string[] lines = Encoding.Default.GetString(Outcar.GetData()).Split('\n');

			for (int i = 0; i < lines.Length; i++)
			{
				string line = lines[i];

				if (line.StartsWith(" POTCAR:") && !potcar.Contains(line))
				{
					potcar.Add(line);
				}

				if (line.Contains("E-fermi"))
				{
					fermiEnergy = double.Parse(line.Split(':')[1].Trim().Split(' ')[0], CultureInfo.InvariantCulture);
				}

				if (line.Contains("NELECT ="))
				{
					electrons = double.Parse(line.Split('=')[1].Trim().Split(' ')[0], CultureInfo.InvariantCulture);
				}

				if (line.Contains("Analysis of symmetry for initial positions (statically)"))
				{
					for (int j = i + 2; j < lines.Length; j++)
					{
						if (lines[j].StartsWith("==="))
						{
							// Next section.
							break;
						}

						if (lines[j].Contains("A1 ="))
						{
							// Extract lattice vectors of the primitive cell.
							SpaceVector[] a = new SpaceVector[3];
							for (int k = 0; k < 3; k++)
							{
								// Assuming format: "A1 = (   2.3216904700,  -4.0212858535,   0.0000000000)"
								string s = lines[j + k];
								a[k] = AtomicEnvironment.ParseSpaceVector(s.Substring(s.IndexOf("(") + 1, s.IndexOf(")") - s.IndexOf("(") - 1));
							}
							primitiveLattice = Lattice.Parse(a[0], a[1], a[2]);
						}

						if (lines[j].Contains("primitive cells build up your supercell"))
						{
							primitiveCells = int.Parse(lines[j].Trim().Split(' ')[0]);
						}

						if (lines[j].Contains("Original cell was already a primitive cell"))
						{
							primitiveCells = 1;
						}
					}
				}

				if (line.Contains("external pressure ="))
				{
					// Expecting unit kbar (kB). Replace with separator. Convert to eV/Å^3.
					externalPressure = double.Parse(line.Replace("kB", "=").Split('=')[1].Trim(), CultureInfo.InvariantCulture) / 1602.176487;
				}
			}

			foreach (string line in potcar)
			{
				string symbol = line.Replace(" POTCAR:", "").Trim().Split(' ')[1];
				((List<Atom>)atoms).Add(Atom.Parse(symbol));
			}
		}
		
		private void ReadPoscar(IEnumerable<Atom> atoms)
		{
			using (MemoryStream stream = new MemoryStream(Poscar.GetData()))
			{
				using (StreamReader reader = new StreamReader(stream))
				{
					initialStructure = VaspPoscar.Read(reader, atoms);
				}
			}
		}
	
		private void ReadVasprun(XmlDocument xml)
		{
			Parse(xml["modeling"]);
		}

		private void ReadOszicar()
		{
			string line = null;
			foreach (string s in Encoding.Default.GetString(Oszicar.GetData()).Split('\n'))
			{
				if (s.Contains("Linear response reoptimize wavefunctions"))
				{
					// Linear response problem. Skip from here on.
					break;
				}

				if (s != "")
				{
					// Use last instance of the "E0=" expression.
					line = s;
				}
			}

			// Read the magnetic moment from OSZICAR, as it doesn't seem to be listed in vasprun.xml.
			if (line.Contains("mag="))
			{
				magneticMoment = double.Parse(line.Substring(line.IndexOf("mag=") + 4).TrimStart(' ').Split(' ')[0], CultureInfo.InvariantCulture);
			}
			else
			{
				// Not spin-polarized calculation.
				magneticMoment = 0.0;
			}

			/*if (readEnergy)
			{
				energy = double.Parse(line.Substring(line.IndexOf("E0=") + 3).TrimStart(' ').Split(' ')[0], CultureInfo.InvariantCulture);
			}*/
		}

		private void ReadElectronicStateDensity()
		{
			using (MemoryStream stream = new MemoryStream(Doscar.GetData()))
			{
				using (StreamReader reader = new StreamReader(stream))
				{
					//electronicStateDensity = VaspElectronicStateDensity.Parse(fermiEnergy, finalStructure.Sites.Count, reader);
					electronicStateDensity = VaspElectronicStateDensity.Parse(fermiEnergy, electrons, reader);
				}
			}
		}

		private void BuildStructure(XmlDocument xml, Lattice primitiveLattice, int primitiveCells, IEnumerable<Atom> atoms)
		{
			// Higher precision in CONTCAR.
			Structure structure;
			using (MemoryStream stream = new MemoryStream(Contcar.GetData()))
			{
				using (StreamReader reader = new StreamReader(stream))
				{
					structure = VaspPoscar.Read(reader, atoms);
				}
			}

			// Try parsing the primitive structure from the primitive basis.
			Structure primitiveStructure = BuildPrimitiveStructure(structure, primitiveLattice, primitiveCells);

			if (primitiveStructure == null && xml != null)
			{
				foreach (string name in new string[] { "finalpos", "initialpos" })
				{
					// Determining primitive structure failed. This may happen for finite difference linear response problems, where
					// CONTCAR containts displaced atoms. In this case try reading from vasprun.xml and reparse primitive structure.
					Structure structure2 = ParseStructure(xml["modeling"].FindNode("structure", name), xml["modeling"]["atominfo"]);

					primitiveStructure = BuildPrimitiveStructure(structure2, primitiveLattice, primitiveCells);
					if (primitiveStructure != null)
					{
						structure = structure2;
						break;
					}
				}
			}

			finalStructure = new VaspStructure(structure.Basis, structure.Lattice, structure.Sites, primitiveStructure);
		}

		private Structure ParseStructure(XmlNode structureNode, XmlNode atomNode)
		{
			Lattice lattice = ParseLattice(structureNode["crystal"].FindNode("varray", "basis"));
		
			List<Atom> atoms = new List<Atom>();
			foreach (XmlNode node in atomNode.FindNode("array", "atomtypes")["set"])
			{
				int n = int.Parse(node.ChildNodes[0].InnerText);
				Atom atom = Atom.Parse(node.ChildNodes[1].InnerText.Trim());
				for (int i = 0; i < n; i++)
				{
					atoms.Add(atom);
				}
			}

			List<Site> sites = new List<Site>();
			foreach (XmlNode node in structureNode.FindNode("varray", "positions"))
			{
				sites.Add(new Site(atoms[0], new Position(lattice, node.ToVector3())));
				atoms.RemoveAt(0);
			}

			if (atoms.Count != 0)
			{
				throw new Exception();
			}

			return new Structure(lattice, sites);
		}

		private Lattice ParseLattice(XmlNode basisNode)
		{
			return Lattice.Parse(basisNode.ChildNodes[0].ToVector3(), basisNode.ChildNodes[1].ToVector3(), basisNode.ChildNodes[2].ToVector3());
		}

		private Structure BuildPrimitiveStructure(Structure structure, Lattice primitiveLattice, int primitiveCells)
		{
			if (primitiveCells == 1)
			{
				// Use exactly same basis.
				primitiveLattice = structure.Lattice;
			}

			if (primitiveCells > 0 && primitiveLattice != null)
			{
				// Tolerance of identical positions in physical distance.
				double tol = 1.0e-4;

				//if (primitiveBasis == null)
				//{
				//	// Couldn't determine the basis. Just ignore it for now.
				//	primitiveBasis = Basis.StandardBasis;
				//}

				//Lattice primitiveLattice = primitiveBasis;
				//Lattice lattice = structure.Lattice;

				List<Site> primitiveSites = new List<Site>();

				foreach (Site site in structure.Sites)
				{
					Atom atom = site.Atom;

					// Real physical position.
					//Vector3 p = lattice.Map * site.Position;
					SpaceVector p = site.Position.PhysicalCoordinates;

					bool found = false;
					foreach (Site site0 in primitiveSites)
					{
						if (site0.Atom != atom)
						{
							continue;
						}

						// Physical position of other atom.
						//Vector3 p0 = lattice.Map * site0.Position;
						SpaceVector p0 = site0.Position.PhysicalCoordinates;

						// Difference relative to the primitive lattice vectors (integer number if identical positions).
						SpaceVector d = primitiveLattice.Decompose(new Position(p - p0)).Coordinates;

						// If the distance in real physical position is below this, then the site is considered the same.
						if (SpaceVector.Norm(primitiveLattice.PhysicalMap * new SpaceVector(d[0] - Math.Round(d[0]), d[1] - Math.Round(d[1]), d[2] - Math.Round(d[2]))) < tol)
						{
							found = true;
							break;
						}
					}

					if (!found)
					{
						primitiveSites.Add(site);
					}
				}

				// Silently ignore discrepancy (don't want to fail if not used), but set to null; only define primitive structure if it matches.
				if (primitiveSites.Count * primitiveCells == structure.Sites.Count)
				{
					return new Structure(primitiveLattice, primitiveSites);
				}
			}

			// Failed. Just return a null reference.
			return null;
		}
		
		private void Parse(XmlNode modelingNode)
		{
			//FinalStructure = ParseStructure(modelingNode.FindNode("structure", "finalpos"), modelingNode["atominfo"]);
			energy = modelingNode["calculation"]["energy"].FindNode("i", "e_wo_entrp").ToDouble();
		}

		public double Energy
		{
			get
			{
				Parse();
				return energy;
			}
		}

		public double MagneticMoment
		{
			get
			{
				Parse();
				return magneticMoment;			
			}
		}

		public double FermiEnergy
		{
			get
			{
				Parse();
				return fermiEnergy;
			}
		}
		
		public double Electrons
		{
			get
			{
				Parse();
				return electrons;
			}
		}

		public double ExternalPressure
		{
			get
			{
				Parse();
				return externalPressure;
			}
		}

		public Structure InitialStructure
		{
			get
			{
				Parse();
				return initialStructure;
			}
		}

		public VaspStructure FinalStructure
		{
			get
			{
				Parse();
				return finalStructure;
			}
		}

		/// <summary>
		/// Electronic density of states (DOSCAR file).
		/// </summary>
		public ElectronicStateDensity ElectronicStateDensity
		{
			get
			{
				ParseElectronicStateDensity();
				return electronicStateDensity;
			}
		}
	}
}
